# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import ujson as json
import functools
import zlib

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.korbit_v1.book_builder import BookBuilder
from coin.exchange.korbit_v1.kr_rest.product import KorbitProduct


@functools.lru_cache(maxsize=64)
def get_product_from_event(event):
  symbol, _ = event.rsplit('_', 1)
  return KorbitProduct.FromStrNativeProduct(symbol)


def decompress_data(data):
  data = data['data']
  data = data.split('|')[1]
  compressed_str = json.loads(data)[0]
  return json.loads(zlib.decompress(bytes.fromhex(compressed_str), 15 + 32))


class KorbitHandler(WebsocketHandler):
  def on_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    product = get_product_from_event(data['event'])
    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return

    book_data = decompress_data(data)
    assert isinstance(book_data, dict)
    bids = book_data['bids']
    asks = book_data['asks']
    book = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self._books[native_symbol] = book
    self.ready = True
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)
    book.reset_observer()

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    product = get_product_from_event(data['event'])
    native_symbol = product.native_symbol
    if not self.is_subscribed('trade', native_symbol):
      return

    trade_data = decompress_data(data)
    assert isinstance(trade_data, list)
    for item in trade_data:
      price = float(item['price'])
      qty = float(item['size'])
      if item['up_or_down'] == 'up':
        side = TradeSide.TRADE_BUY_SIDE
      elif item['up_or_down'] == 'down':
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError('Unknown side: %s', trade_data)
      trade_event = create_market_trade_event(record.timestamp,
                                              native_symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', native_symbol, trade_event)
